#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <stdlib.h>
#include <errno.h>
#include <stdio.h>
#include <string.h>
#include <poll.h>   // poll function and structs
#include <limits.h> // For OPEN_MAX (alternative definition)

#define MAXLINE 1024

#ifndef OPEN_MAX
#define OPEN_MAX 1024
#endif

int main(int argc, char **argv)
{
    int servPort = 6888; // Port from PDF example
    int listenq = 1024;  // Listen queue size
    int listenfd, connfd;
    struct pollfd clients[OPEN_MAX];
    int max_clients_idx = 0; // Tracks the highest index used in clients array

    socklen_t socklen = sizeof(struct sockaddr_in);
    struct sockaddr_in cliaddr, servaddr;
    char buf[MAXLINE];
    int nready; // Number of descriptors ready from poll()

    listenfd = socket(AF_INET, SOCK_STREAM, 0);
    if (listenfd < 0)
    {
        perror("socket error");
        exit(1);
    }

    int opt = 1;
    if (setsockopt(listenfd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0)
    {
        perror("setsockopt error");
        close(listenfd);
        exit(1);
    }

    memset(&servaddr, 0, sizeof(servaddr));
    servaddr.sin_family = AF_INET;
    servaddr.sin_addr.s_addr = INADDR_ANY;
    servaddr.sin_port = htons(servPort);

    if (bind(listenfd, (struct sockaddr *)&servaddr, socklen) == -1)
    {
        perror("bind error");
        close(listenfd);
        exit(1); // Use consistent exit code
    }

    if (listen(listenfd, listenq) < 0)
    {
        perror("listen error");
        close(listenfd);
        exit(1);
    }

    clients[0].fd = listenfd;
    clients[0].events = POLLIN;
    clients[0].revents = 0;

    for (int i = 1; i < OPEN_MAX; i++)
    {
        clients[i].fd = -1;
        clients[i].events = 0;
        clients[i].revents = 0;
    }
    max_clients_idx = 0;

    printf("poll echo server startup, listen on port: %d\n", servPort);
    printf("max connections (array size): %d\n", OPEN_MAX);

    for (;;)
    {
        nready = poll(clients, max_clients_idx + 1, -1);

        if (nready == -1)
        {
            if (errno == EINTR)
            {
                continue;
            }
            perror("poll error");
            continue;
        }

        if (clients[0].revents & POLLIN)
        {
            connfd = accept(listenfd, (struct sockaddr *)&cliaddr, &socklen);
            if (connfd < 0)
            {
                if (errno != EINTR && errno != EWOULDBLOCK && errno != EAGAIN)
                {
                    perror("accept error");
                }
            }
            else
            {
                sprintf(buf, "accept form %s:%d\n", inet_ntoa(cliaddr.sin_addr), ntohs(cliaddr.sin_port));
                printf("%s", buf);

                int i;
                for (i = 1; i < OPEN_MAX; i++)
                {
                    if (clients[i].fd < 0)
                    {
                        clients[i].fd = connfd;
                        clients[i].events = POLLIN;
                        clients[i].revents = 0;
                        if (i > max_clients_idx)
                        {
                            max_clients_idx = i;
                        }
                        break;
                    }
                }

                if (i == OPEN_MAX)
                {
                    fprintf(stderr, "too many connections, more than %d\n", OPEN_MAX - 1);
                    close(connfd);
                }
            }

            if (--nready <= 0)
            {
                continue;
            }
        }

        for (int i = 1; i <= max_clients_idx; i++)
        {
            int sockfd = clients[i].fd;
            if (sockfd < 0)
            {
                continue;
            }

            if (clients[i].revents & (POLLIN | POLLERR))
            {
                ssize_t nread = read(sockfd, buf, MAXLINE - 1);

                if (nread < 0)
                {
                    if (errno == ECONNRESET)
                    {
                        printf("Client fd %d connection reset.\n", sockfd);
                    }
                    else
                    {
                        perror("read error");
                    }
                    close(sockfd);
                    clients[i].fd = -1;
                }
                else if (nread == 0)
                {
                    printf("Client fd %d closed the connection.\n", sockfd);
                    close(sockfd);
                    clients[i].fd = -1;
                }
                else
                {
                    buf[nread] = '\0';
                    printf("Server received from fd %d: %s\n", sockfd, buf);
                }

                if (--nready <= 0)
                {
                    break;
                }
            }
        }
    }

    close(listenfd);
    return 0;
}